'''
    The Function of SearchImage(Rubbish English -- something wrong with my input method)
'''

import h5py
import matplotlib.image as mpimg
import matplotlib.pyplot as plt
import numpy as np
from Net.VGG16 import SE_VGG
import LoadImage as lim
import torchvision.models as models
import torch

query = 'goal.jpg' # use this to search
index = 'FeaturesAndPaths/vgg_featureCNN.h5' # location
result = 'DB/pictures'

h5f = h5py.File(index,'r') # open this file.h5
features = h5f['features'][:] # get features
paths = h5f['paths'][:] # get paths
h5f.close() # close this file

# print(features)
# print(paths)

print('------------------------------------------\n'
      '            search image start            \n'
      '------------------------------------------')

# read and show query image
queryImage = mpimg.imread(query)
plt.title("The image I want to find")
plt.imshow(queryImage)
plt.show()

# init VGG16 model (custom dimensionality)
model = models.vgg16(pretrained=False)
pthfile = 'Model/state/vgg16-397923af.pth'  # 预训练模型的地址
model.load_state_dict(torch.load(pthfile))
model.eval()  # 表示用于预测

# # Function
# def searchImageFun():
#     img = lim.synthesize(query)
#     queryFeat = model(img)
#     queryFeat = queryFeat.view(queryFeat.size(1))  # transform to one-dimensional
#     # use cosine similarity method
#     scores = np.dot(queryFeat.detach().numpy(), features.T)  # inner product
#     # print(scores)
#     rank_ID = np.argsort(scores)[::-1]  # sort and return index value
#     # print(rank_ID)
#     rank_scores = scores[rank_ID]
#     # print(rank_scores)
#
#     maxres = 3  # number that need to be return
#     imlist = []  # store the most likely pictures
#     for epoch_idx, index in enumerate(rank_ID[0:maxres]):
#         imlist.append(paths[index])
#     return imlist
#
if __name__ == '__main__':
    img = lim.synthesize(query)
    queryFeat = model(img)
    queryFeat = queryFeat.view(queryFeat.size(1)) # transform to one-dimensional
    # use cosine similarity method
    scores = np.dot(queryFeat.detach().numpy(),features.T) # inner product
    # print(scores)
    rank_ID = np.argsort(scores)[::-1] # sort and return index value
    # print(rank_ID)
    rank_scores = scores[rank_ID]
    # print(rank_scores)

    maxres = 5 # number that need to be return

    imlist = [] # store the most likely pictures
    for epoch_idx, index in enumerate(rank_ID[0:maxres]):
        imlist.append(paths[index])
        print(f"imagePath : {str(paths[index])} scores = {rank_scores[epoch_idx]}")
        print("top %d images in order are: " %maxres , imlist)

    for idx,path in enumerate(imlist):
        print(path)
        image = mpimg.imread(str(path,'utf-8')) # get picture
        plt.title('Searching Results')
        plt.imshow(image)
        plt.show()


